/*
 *  Copyright 2008-2009 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */


#include <algorithm>
#include <vector>
#include <stdio.h>

#include <cusp/blas.h>
#include <cusp/csr_matrix.h>

extern "C" {
#include <metis.h>
    extern int METIS_PartGraphKway(idx_t *nvtxs, idx_t *ncon, idx_t *xadj,
                                   idx_t *adjncy, idx_t *vwgt, idx_t *vsize, idx_t *adjwgt,
                                   idx_t *nparts, real_t *tpwgts, real_t *ubvec, idx_t *options,
                                   idx_t *objval, idx_t *part);

    extern int METIS_PartGraphRecursive(idx_t *nvtxs, idx_t *ncon, idx_t *xadj,
                                        idx_t *adjncy, idx_t *vwgt, idx_t *vsize, idx_t *adjwgt,
                                        idx_t *nparts, real_t *tpwgts, real_t *ubvec, idx_t *options,
                                        idx_t *objval, idx_t *part);
};



namespace cusp
{
namespace graph
{
namespace detail
{
namespace host
{
template <typename MethodType, class MatrixType, class ArrayType>
size_t partition(const MatrixType& graph,
                 size_t num_parts,
                 ArrayType& part,
                 bool Kway,
                 cusp::graph::metis)
{
    typedef typename MatrixType::index_type IndexType;
    typedef typename MatrixType::value_type ValueType;

    if (graph.num_rows != graph.num_cols)
        throw cusp::invalid_input_exception("Input graph must be structurally symmetric");

    if (num_parts < 2) {
        // only one partition
        cusp::blas::fill(part, 0);
        return 0;
    }

    part.resize(graph.num_rows);
    ArrayType int_partition(part.size()); //because METIS expects an int array

    idx_t options[METIS_NOPTIONS];
    METIS_SetDefaultOptions(options);
    idx_t ncon = 1, edgecut;

    printf(" [A=A^T:");
    bool is_symmetric = cusp::graph::is_structurally_symmetric(graph);
    if (is_symmetric)
        printf(" yes]");
    else
        printf(" no]");


    if (is_symmetric) {
        if (Kway)
            METIS_PartGraphKway((IndexType *) &graph.num_rows,
                                (IndexType *) &ncon,
                                (IndexType *) thrust::raw_pointer_cast(&graph.row_offsets[0]),
                                (IndexType *) thrust::raw_pointer_cast(&graph.column_indices[0]),
                                (IndexType *) NULL, (IndexType *) NULL, (IndexType *) NULL,
                                &num_parts,
                                (real_t *) NULL, (real_t *) NULL,
                                options, &edgecut,
                                (IndexType *) thrust::raw_pointer_cast(&int_partition[0]));
        else
            METIS_PartGraphRecursive((IndexType *) &graph.num_rows,
                                     (IndexType *) &ncon,
                                     (IndexType *) thrust::raw_pointer_cast(&graph.row_offsets[0]),
                                     (IndexType *) thrust::raw_pointer_cast(&graph.column_indices[0]),
                                     (IndexType *) NULL, (IndexType *) NULL, (IndexType *) NULL,
                                     &num_parts,
                                     (real_t *) NULL, (real_t *) NULL,
                                     options, &edgecut,
                                     (IndexType *) thrust::raw_pointer_cast(&int_partition[0]));
    } else {
        //printf("Matrix graph is nonsymmetric, using A^T + A instead\n");

        cusp::coo_matrix<IndexType,ValueType,cusp::host_memory> coo_AtA( graph );

        cusp::blas::fill(coo_AtA.values,ValueType(1));

        MatrixType csr_AtA( graph );

        assert(cusp::graph::is_structurally_symmetric(csr_AtA));

        if (Kway)
            METIS_PartGraphKway((IndexType *) &csr_AtA.num_rows,
                                (IndexType *) &ncon,
                                (IndexType *) thrust::raw_pointer_cast(&csr_AtA.row_offsets[0]),
                                (IndexType *) thrust::raw_pointer_cast(&csr_AtA.column_indices[0]),
                                (IndexType *) NULL, (IndexType *) NULL, (IndexType *) NULL,
                                &num_parts,
                                (real_t *) NULL, (real_t *) NULL,
                                options, &edgecut,
                                (IndexType *) thrust::raw_pointer_cast(&int_partition[0]));
        else
            METIS_PartGraphRecursive((IndexType *)&csr_AtA.num_rows,
                                     (IndexType *) &ncon,
                                     (IndexType *) thrust::raw_pointer_cast(&csr_AtA.row_offsets[0]),
                                     (IndexType *) thrust::raw_pointer_cast(&csr_AtA.column_indices[0]),
                                     (IndexType *) NULL, (IndexType *) NULL, (IndexType *) NULL,
                                     &num_parts,
                                     (real_t *) NULL, (real_t *) NULL,
                                     options, &edgecut,
                                     (IndexType *) thrust::raw_pointer_cast(&int_partition[0]));
    }

    printf(" [edgecut: %4.1f%%]\n", 100* (double) edgecut/ (double) graph.num_entries);

    //copy partition to vector
    cusp::blas::copy(int_partition, part);

    return (IndexType) edgecut;
}

} // end namespace host
} // end namespace detail
} // end namespace graph
} // end namespace cusp

